import random
from datetime import datetime
import numpy as nm
import scipy.sparse as sp
from datetime import datetime
from logging import getLogger
from recbole.utils import set_color


# 包装成一个函数，叫adjacent_random_augment,输入输出都是类似liked_graph那样的邻接矩阵
# 注意函数里自带两个变量A和B，是(0, 1)的参数，先定义好，A=0.4，B=0.5
# 函数的输入是liked_graph
# 函数的任务是统计liked_graph里零元素的数量，在liked_graph里选比例为A的零元素，具体数量向下取整
# 然后把这些元素的值从0变成B，得到的结果存在一个新矩阵
# 输出得到的新矩阵变量，叫random_augmented_like_graph
#  A即sample_proportion，B即augment_value
def random_graph_augment(liked_graph, sample_proportion=0.1, augment_value=1):
    # 复制一份基础的交互矩阵
    augmented_liked_graph = liked_graph.copy()

    # 找出0元素的id
    zero_index = nm.argwhere(augmented_liked_graph == 0)
    num_of_zero = len(zero_index)

    # 根据预定比例确定随机增强的元素个数
    sample_num = int(num_of_zero * sample_proportion)
    # 以及随机增强元素的id
    random_index = random.sample(zero_index.tolist(), sample_num)

    # 对指定0元素进行增强赋值
    for index in random_index:
        augmented_liked_graph[index[0], index[1]] = augment_value

    return augmented_liked_graph

# 基于sp.find执行零元素坐标对采样的随机增强
def find_random_matrix_augment(sparse_matrix, sample_proportion=0.1, augment_value=1):
    # 获得lil形式的副本
    augment_matrix = sp.lil_matrix(sparse_matrix)
    # 获得零元素id与数量
    zero_tuple = sp.find(augment_matrix == 0)
    zero_index = list(zip(zero_tuple[0], zero_tuple[1]))
    num_of_zero = len(zero_index)
    # 随机确定要增强的元素
    sample_num = int(num_of_zero * sample_proportion)
    random_index = random.sample(zero_index, sample_num)
    # 增强赋值
    for index in random_index:
        augment_matrix[index[0], index[1]] = augment_value
    return augment_matrix


# 基于零元素采样列表维护的随机增强，采样数量基于非零元素量
def sample_random_matrix_augment(sparse_matrix, sample_proportion=0.1, augment_value=1):
    logger = getLogger()
    # print(datetime.now(), ": 开始large_random_matrix_augment")
    logger.info(f'{datetime.now()}' + set_color(': sample_random_matrix_augment begin', 'green'))

    # 获得lil形式的副本
    augment_matrix = sp.lil_matrix(sparse_matrix)

    # 随机确定要增强的元素
    sample_num = int(sparse_matrix.nnz * sample_proportion)
    random_index = sample_zero_n(sparse_matrix, sample_num)

    # 增强赋值
    for index in random_index:
        augment_matrix[index[0], index[1]] = augment_value
    return augment_matrix

# 基于零元素采样的随机增强，对高稀疏度矩阵作了优化，采样数量基于非零元素量
def large_random_matrix_augment(sparse_matrix, sample_proportion=0.1, augment_value=1):
    logger = getLogger()
    # print(datetime.now(), ": 开始large_random_matrix_augment")
    logger.info(f'{datetime.now()}' + set_color(': large_random_matrix_augment begin', 'green'))
    # 获得lil形式的副本
    augment_matrix = sp.lil_matrix(sparse_matrix)

    # 随机确定要增强的元素，并立即赋值
    sample_num = int(sparse_matrix.nnz * sample_proportion)
    for _ in range(sample_num):
        sampled_index = nm.random.randint(0, [sparse_matrix.shape[0], sparse_matrix.shape[1]])
        augment_matrix[sampled_index[0], sampled_index[1]] = augment_value

    return augment_matrix


def sample_zero_n(mat, n):
    """
    :param mat: 原稀疏矩阵
    :param n:   要取的零元素数量
    :return:    取到的0元素坐标列表
    """
    itr = sample_zero_forever(mat)
    return [next(itr) for _ in range(n)]


# sample_zero_n中使用的迭代器。维护一个列表，由非0元素与已取过的0元素共同组成
# ，每次取0元素时在列表中判断是否重复，不重复则取到，yield返回，并加入列表
def sample_zero_forever(mat):
    nonzero_or_sampled = set(zip(*mat.nonzero()))
    while True:
        t = tuple(nm.random.randint(0, [mat.shape[0], mat.shape[1]]))
        if t not in nonzero_or_sampled:
            yield t
            nonzero_or_sampled.add(t)

